from MCP协议与AIagent开发应用.第6章MCP应用开发进阶.子任务嵌套 import Prompt, ContextChain


class MailModuleBase:
    def __init__(self,mmodule_name:str):
        self.mmodule_name = mmodule_name

    def inject_params(self,params:dict):
        raise NotImplementedError


    def process(self) -> str:
        raise NotImplementedError


class MailSorter(MailModuleBase):
    """
    邮件整理模块
    """
    def __init__(self):
        super().__init__("邮件整理")
        self.mail_content = ""
        self.mail_subject = ""

    def inject_params(self, params: dict):
            self.mail_subject = params.get("subject", "")
            self.mail_content = params.get("content", "")

    def process(self) -> str:
            # 简单逻辑：若主题或内容包含"紧急"则分类为紧急，否则普通
            if "紧急" in self.mail_subject or "紧急" in self.mail_content:
                return f"邮件分类: [紧急]\n主题: {self.mail_subject}"
            else:
                return f"邮件分类: [普通]\n主题: {self.mail_subject}"

class MailFilter(MailModuleBase):
        """
        邮件过滤模块，筛选不需要的邮件，如广告、垃圾邮件等。
        """

        def __init__(self):
            super().__init__("邮件过滤")
            self.classification = ""
            self.mail_content = ""

        def inject_params(self, params: dict):
            self.classification = params.get("classification", "")
            self.mail_content = params.get("content", "")

        def process(self) -> str:
            # 根据分类，若含有敏感词"广告"则过滤，否则通过
            if "广告" in self.mail_content or "广告" in self.classification:
                return "结果: 该邮件被过滤 (广告邮件)"
            return f"结果: 该邮件通过过滤 ({self.classification})"

class MailSender(MailModuleBase):
        """
        邮件发送模块，负责将结果发送给相应收件人。
        """

        def __init__(self):
            super().__init__("邮件发送")
            self.to = ""
            self.body = ""

        def inject_params(self, params: dict):
            self.to = params.get("to", "")
            self.body = params.get("body", "")

        def process(self) -> str:
            return f"邮件已发送给 {self.to}:\n{self.body}"

class MailArchiver(MailModuleBase):
        """
        邮件归档模块，负责将已处理过的邮件存储于归档仓库，供后续查询。
        """

        def __init__(self):
            super().__init__("邮件归档")
            self.archive_path = ""
            self.log_info = ""

        def inject_params(self, params: dict):
            self.archive_path = params.get("archive_path", "/default/archive/")
            self.log_info = params.get("log_info", "")

        def process(self) -> str:
            return f"邮件信息已归档至 {self.archive_path}\n附加日志: {self.log_info}"

    ################################################################################
    # 场景执行：邮件整理 + 邮件过滤 + 邮件发送 + 邮件归档
    ################################################################################

def run_mail_pipeline():
        """
        演示基于MCP协议思想的可复用任务模块与参数注入流程。
        """

        # 1. 构造上下文管理器
        chain = ContextChain()

        # 2. 添加系统Prompt，描述流程要求
        system_prompt = Prompt(
            role="system",
            content="这是一个邮件处理自动化系统，支持整理、过滤、发送与归档操作。"
        )
        chain.add_prompt(system_prompt)

        # 3. 用户输入，提供邮件主题与内容
        user_prompt = Prompt(
            role="user",
            content="邮件主题: '紧急会议通知'\n邮件正文: '这是一份需要即时处理的紧急会议。'"
        )
        chain.add_prompt(user_prompt)

        # ============ Step A: 邮件整理模块 ============
        sorter = MailSorter()
        sorter.inject_params({
            "subject": "紧急会议通知",
            "content": "这是一份需要即时处理的紧急会议。"
        })
        sort_result = sorter.process()

        # 将处理结果注入上下文
        sorter_prompt = Prompt(
            role="assistant",
            content=sort_result,
            parent_id=user_prompt.id
        )
        chain.add_prompt(sorter_prompt)

        # ============ Step B: 邮件过滤模块 ============
        filter_mod = MailFilter()
        filter_mod.inject_params({
            "classification": "紧急",
            "content": "这是一份需要即时处理的紧急会议。"
        })
        filter_result = filter_mod.process()

        filter_prompt = Prompt(
            role="assistant",
            content=filter_result,
            parent_id=sorter_prompt.id
        )
        chain.add_prompt(filter_prompt)

        # ============ Step C: 邮件发送模块 ============
        # 若通过过滤，才进入发送环节
        if "通过过滤" in filter_result:
            sender = MailSender()
            sender.inject_params({
                "to": "boss@example.com",
                "body": "紧急会议邮件: 需要即时处理，请查看内容。"
            })
            send_result = sender.process()

            send_prompt = Prompt(
                role="assistant",
                content=send_result,
                parent_id=filter_prompt.id
            )
            chain.add_prompt(send_prompt)
        else:
            no_send_prompt = Prompt(
                role="assistant",
                content="该邮件未通过过滤，不会进行发送。",
                parent_id=filter_prompt.id
            )
            chain.add_prompt(no_send_prompt)

        # ============ Step D: 邮件归档模块 ============
        # 进行归档，保存处理记录
        archiver = MailArchiver()
        archiver.inject_params({
            "archive_path": "/mail/archive/2025/urgent/",
            "log_info": f"{sort_result}\n{filter_result}"
        })
        archive_result = archiver.process()
        archive_prompt = Prompt(
            role="assistant",
            content=archive_result,
            parent_id=(send_prompt.id if 'send_prompt' in locals() else filter_prompt.id)
        )
        chain.add_prompt(archive_prompt)

        # 4. 打印最终上下文链，查看执行流程
        chain.print_context(archive_prompt.id)

if __name__ == "__main__":
    run_mail_pipeline()